import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np


def plot_loss(loss_value, savepath):
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 3.7))
    ax.plot(loss_value)
    ax.set_title('Loss')
    fig.savefig(savepath)


def plot_cluster_distribution(labels_sequence, rawdata):
    n, m = labels_sequence.shape
    idx = list(np.arange(0, n))
    fig, axes = plt.subplots(nrows=m+1, ncols=1, figsize=(20, 25))
    #print(rawdata.shape)
    sns.lineplot(data=rawdata, ax=axes[0])
    axes[0].set_title('Raw')
    for k in range(m):
         sns.heatmap(data=labels_sequence[:, k].reshape(1, -1),  ax=axes[k+1])
    plt.show()
    plt.savefig('/home/kexin/f.eps')


